from models import Swin_TransformerModel, Swin_Model_Without_Mlp, Swin_Diff_Head_Decode, swin_obj_dect_trans,\
    feature_fusion, feature_fuse_with_cnn_feature, feature_fuse, bert_feature_fuse
import os
import torch


def setup(opt):
    if opt.caption_model == 'swin_transformer':
        model = Swin_TransformerModel.SwinTransformerModel(opt)
    elif opt.caption_model == 'swin_without_mlp':
        model = Swin_Model_Without_Mlp.SwinModel(opt)
    elif opt.caption_model == 'swin_diff_head_decode':
        model = Swin_Diff_Head_Decode.SwinDifferentHead(opt)
    elif opt.caption_model == 'swin_obj_dect_trans':
        model = swin_obj_dect_trans.Swin_obj_dect_trans(opt)
    elif opt.caption_model == 'feature_fusion':
        model = feature_fusion.Feature_fusion(opt)
    elif opt.caption_model == 'feature_fusion2':
        model = feature_fuse_with_cnn_feature.Feature_fusion(opt)
    elif opt.caption_model == 'feature_fusion3':
        model = feature_fuse.Feature_fusion(opt)
    elif opt.caption_model == 'bert_feature_fuse':
        model = bert_feature_fuse.Bert_Feature_Fuse(opt)

    if vars(opt).get('start_from', None) is not None:
        assert os.path.isdir(opt.start_from), " %s must be a path" % opt.start_from
        assert os.path.isfile(os.path.join(opt.start_from, 'infos_' + opt.id + ".pkl")), "infos.pkl file does not exist in path %s" % opt.start_from
        model.load_state_dict(torch.load(os.path.join(opt.start_from, 'model.pth')))

    return model